package src;

public class KeyExpansion {
    private static void word_loop(byte[] in, byte[] out){
        for(int i = 0; i < 4; i++){
            out[i] = in[(i + 1) % 4];
        }
    }

    private static void byte_substitution(byte[] in){
        for(int i = 0; i < 4; i++){
            in[i] = (byte)Const.s_box[Byte.toUnsignedInt(in[i])];
        }
    }

    private static void round_xor(byte[] byte_substitution_out, int round_num){
        byte_substitution_out[0] = (byte)((int)byte_substitution_out[0] ^ Const.round_constant[round_num]);
    }

    public static void key_expansion(byte[] matrix){
        byte[] in = new byte[4];
        byte[] out = new byte[4];

        for(int i = 4; i < 44; i++){
            int round_num = (i / 4) - 1;
            if(i % 4 != 0){//当生成的密钥列不是4的倍数时
                for(int row = 0; row < 4; row++){//直接进行异或处理
                    matrix[i * 4+ row] = (byte) (matrix[(i - 4) * 4+ row] ^ matrix[(i - 1) * 4 + row]);
                }
            }else {//进行T函数处理
                for(int j = 0; j < 4; j++){//将第i - 1 列赋值给in
                    in[j] = matrix[(i - 1) * 4 + j];
                }
                word_loop(in, out);
                byte_substitution(out);
                round_xor(out, round_num);
                for(int row = 0; row < 4; row++){
                    matrix[i * 4 + row] = (byte)(matrix[(i - 4) * 4 + row] ^ out[row]);
                }
            }
        }
    }
}
